#include "utils/object_detector.hpp"


ObjectDetector::ObjectDetector (InferenceEngine::Core &ie,
                            const std::string &device,
                            const std::string &model)
{
    auto network = ie.ReadNetwork(model);
    auto input = network.getInputsInfo();
    auto inputIter = input.begin();
    inputName = inputIter->first;
    auto inputInfo = inputIter->second;
    inputInfo->setPrecision(InferenceEngine::Precision::U8);
    inputInfo->setLayout(InferenceEngine::Layout::NCHW);
    auto output = network.getOutputsInfo();
    auto outputIter = output.begin();
    outputName = outputIter->first;
    auto outputInfo = outputIter->second;
    outputInfo->setPrecision(InferenceEngine::Precision::FP32);
    net = ie.LoadNetwork(network, device);
}

auto ObjectDetector::infer (const cv::Mat &frame, float confidence) -> std::vector<Result> {
    auto req = net.CreateInferRequest();
    auto inputBlob = req.GetBlob(inputName);
    input(frame, inputBlob);
    req.Infer();
    std::vector<Result> results;
    auto outputBlob = req.GetBlob(outputName);
    output(outputBlob, results, confidence);
    return results;
}

void ObjectDetector::input (const cv::Mat &src, InferenceEngine::Blob::Ptr &dst) {
    auto size = dst->getTensorDesc().getDims();
    auto w = (int)size[3];
    auto h = (int)size[2];
    auto c = (int)size[1];
    cv::Mat normalized(src);
    cv::resize(src, normalized, cv::Size(w, h));
    auto imap = InferenceEngine::as<InferenceEngine::MemoryBlob>(dst)->wmap();
    auto ip = imap.as<uint8_t*>();
    if (c == 1) {
        typedef float f32v4 __attribute__((vector_size(16)));
        f32v4 gray = {0.114f, 0.587f, 0.299f, 0.0f};
        for (auto y = 0; y < h; ++y) {
            auto rp = ip + w * y;
            for (auto x = 0; x < w; ++x) {
                auto *s = (uint32_t*)normalized.at<cv::Vec3b>(y, x).val;
                uint32_t d;
                f32v4 t;
                asm volatile (R"(
                    VPMOVZXBD (%[s]), %[t]
                    VCVTDQ2PS %[t], %[t]
                    VDPPS $0x71, %[f], %[t], %[t]
                    VCVTSS2SI %[t], %[d]
                )"
                :[d]"=r"(d), [t]"=&x"(t)
                :[s]"r"(s), [f]"x"(gray)
                );
                rp[x] = (uint8_t)d;
            }
        }
    } else if (c == 3) {
        for (auto z = 0; z < c; ++z) {
            auto cp = ip + w * h * z;
            for (auto y = 0; y < h; ++y) {
                auto rp = cp + w * y;
                for (auto x = 0; x < w; ++x)
                    rp[x] = normalized.at<cv::Vec3b>(y, x)[z];
            }
        }
    } else {
        throw std::logic_error("");
    }
}

void ObjectDetector::output (const InferenceEngine::Blob::Ptr &src, std::vector<Result> &dst, float confidence) {
    auto op = InferenceEngine::as<InferenceEngine::MemoryBlob>(src)->rmap().as<float*>();
    auto maxCount = src->getTensorDesc().getDims()[2];
    for (auto i = 0; i < maxCount; ++i) {
        auto p = op + (i * 7);
        auto id = (size_t)p[0];
        if (id < 0)
            break;
        auto conf = p[2];
        if (conf < confidence)
            continue;
        auto type = (size_t)p[1];
        auto x = std::max(p[3], 0.0f);
        auto y = std::max(p[4], 0.0f);
        auto w = std::min(p[5], 1.0f) - x;
        auto h = std::min(p[6], 1.0f) - y;
        dst.emplace_back(id, type, conf, x, y, w, h);
    }
}
